import json
import time
import requests
import re

class OllamaAPI:
    def __init__(self, config_file='config.json', history_file='user_messages.json'):
        # 从配置文件中读取配置
        with open(config_file, 'r') as file:
            config = json.load(file)
        
        self.api_url = config.get('api_url', 'http://localhost:11434/api/chat')
        self.model = config.get('model', 'deepseek-r1:7b')

        self.history_file = history_file
        self.user_messages = {}  # 存储用户对话历史

        # 尝试加载历史数据
        self.load_messages()

    def load_messages(self):
        """从文件加载用户对话历史"""
        try:
            with open(self.history_file, 'r') as file:
                self.user_messages = json.load(file)
        except FileNotFoundError:
            self.user_messages = {}  # 如果文件不存在，初始化为空字典

    def save_messages(self):
        """将用户对话历史保存到文件"""
        with open(self.history_file, 'w') as file:
            json.dump(self.user_messages, file, ensure_ascii=False, indent=4)

    def get_reply(self, sender_name: str, prompt: str):
        """根据用户对话历史生成回复"""
        # 获取该用户的对话历史
        messages = self.user_messages.get(sender_name, [])
        
        # 添加当前用户的提问
        messages.append({"role": "user", "content": prompt})

        # 请求数据
        headers = {'Content-Type': 'application/json'}
        data = {
            'model': self.model,
            'messages': messages,
            'stream': False
        }

        response = requests.post(self.api_url, json=data, headers=headers)

        if response.status_code != 200:
            print(f"Error: {response.status_code}")
            return None

        # 解析返回的数据
        response_json = response.json()

        # 获取生成的完整回复内容
        full_response = response_json.get('response', "").strip()

        # 移除 <think> 和 </think> 标签及其中的内容
        full_response = re.sub(r'<think>.*?</think>', '', full_response, flags=re.DOTALL)

        # 将回复内容添加到历史中
        messages.append({"role": "assistant", "content": full_response})
        self.user_messages[sender_name] = messages
        
        # 保存更新后的对话历史
        self.save_messages()

        return full_response.strip()